﻿using System;
using System.Text;
using System.Diagnostics;

namespace GypeInject
{
    public static class Injector
    {
        public static Win32.ModuleEntry32 Inject(Process process, string dllPath, out IntPtr startAddress)
        {
            Win32.ModuleEntry32 outModule = new Win32.ModuleEntry32();
            var processHandle = Win32.OpenProcess(Win32.ProcessAccessRights.ALL_ACCESS, false, process.Id);

            InjectModule(processHandle, dllPath);

            var baseAddress = IntPtr.Zero;
            Win32.EnumProcessModules(process.Id, delegate(Win32.ModuleEntry32 module)
            {
                if (module.ModulePath == dllPath)
                {
                    baseAddress = module.BaseAddress;
                    outModule = module;


                    return true;
                }
                System.Diagnostics.Debug.Write(module.ModulePath);
                System.Diagnostics.Debug.WriteLine(" " + module.BaseAddress.ToInt64());
                return false;
            });
            /*
            foreach (ProcessModule module in process.Modules)
            {
                if (module.FileName == dllPath)
                {
                    baseAddress = module.BaseAddress;
                    outModule = module;
                    break;
                }
            }
            */
            if (baseAddress == IntPtr.Zero)
                throw new Exception("Could not inject!");

            var lib = Win32.LoadLibrary(dllPath);
            var startAddr = Win32.GetProcAddress(lib, "Start");


            if (Environment.Is64BitOperatingSystem)
            {
                var addr = baseAddress.ToInt64() - lib.ToInt64() + startAddr.ToInt64();
                startAddress = new IntPtr(addr);
            }
            else
            {
                var addr = baseAddress.ToInt32() - lib.ToInt32() + startAddr.ToInt32();
                startAddress = new IntPtr(addr);
            }

            Win32.FreeLibrary(lib);

            return outModule;
        }

        private static void InjectModule(IntPtr processHandle, string dllPath)
        {
            int numBytesWritten;
            var address = Win32.VirtualAllocEx(processHandle, IntPtr.Zero, dllPath.Length + 1,
                                Win32.AllocationType.COMMIT, Win32.ProtectionFlags.EXECUTE_READWRITE);

            if (address == IntPtr.Zero)
                throw new Exception("VirtualAllocEx failed!");

            Win32.WriteProcessMemory(processHandle, address, Encoding.ASCII.GetBytes(dllPath), dllPath.Length + 1, out numBytesWritten);

            var moduleHandle = Win32.GetModuleHandle("Kernel32.dll");
            var loadLibAddr = Win32.GetProcAddress(moduleHandle, "LoadLibraryA");

            uint threadId;
            var threadHandle = Win32.CreateRemoteThread(processHandle, IntPtr.Zero, 0, loadLibAddr, address, Win32.CreationFlags.RUN, out threadId);
            Win32.WaitForSingleObject(threadHandle, int.MaxValue);

            Win32.VirtualFreeEx(processHandle, address, 0, Win32.DeallocationType.RELEASE);
        }

        public static void CallStart(Process process, IntPtr address)
        {
            var processHandle = Win32.OpenProcess(Win32.ProcessAccessRights.ALL_ACCESS, false, process.Id);

            uint threadId;
            var handle = Win32.CreateRemoteThread(processHandle, IntPtr.Zero, 0, address, IntPtr.Zero, Win32.CreationFlags.RUN, out threadId);
            Win32.WaitForSingleObject(handle, int.MaxValue);
        }

        public static void UnloadDll(Process process, Win32.ModuleEntry32 module)
        {
            var processHandle = Win32.OpenProcess(Win32.ProcessAccessRights.ALL_ACCESS, false, process.Id);

            var lib = Win32.LoadLibrary(module.ModulePath);
            var endAddress = Win32.GetProcAddress(lib, "End");

#if X64
            var addr = module.BaseAddress.ToInt64() - lib.ToInt64() + endAddress.ToInt64();
#else
            var addr = module.BaseAddress.ToInt32() - lib.ToInt32() + endAddress.ToInt32();
#endif

            uint threadId;
            Win32.CreateRemoteThread(processHandle, IntPtr.Zero, 0, new IntPtr(addr), IntPtr.Zero, Win32.CreationFlags.RUN, out threadId);

            var moduleHandle = Win32.GetModuleHandle("Kernel32.dll");
            var loadLibAddr = Win32.GetProcAddress(moduleHandle, "FreeLibrary");

            var threadHandle = Win32.CreateRemoteThread(processHandle, IntPtr.Zero, 0, loadLibAddr, module.BaseAddress, Win32.CreationFlags.RUN, out threadId);
            Win32.WaitForSingleObject(threadHandle, int.MaxValue);
        }
    }
}
